use std::collections::VecDeque;
use std::mem;

/// A node in the trie, containing a value and an array of child nodes.
struct TrieNode<T> {
    data: Option<T>,
    use_count: usize,
    next: [Option<Box<TrieNode<T>>>; 256],
}

impl<T> TrieNode<T> {
    /// Creates a new `TrieNode` with no data.
    fn new() -> Self {
        TrieNode {
            data: None,
            use_count: 0,
            next: Default::default(),
        }
    }
}

/// A trie data structure for mapping strings to values.
pub struct Trie<T> {
    root_node: Option<Box<TrieNode<T>>>,
}

impl<T> Trie<T> {
    /// Creates a new, empty `Trie`.
    pub fn new() -> Self {
        Trie { root_node: None }
    }

    /// Inserts a key-value pair into the trie.
    /// Returns `true` if the insertion was successful, `false` otherwise.
    pub fn insert(&mut self, key: &str, value: T) -> bool {
        if key.is_empty() {
            return false;
        }

        let mut rover = &mut self.root_node;
        let key_bytes = key.as_bytes();

        for &c in key_bytes {
            let node = rover.get_or_insert_with(|| Box::new(TrieNode::new()));
            node.use_count += 1;
            rover = &mut node.next[c as usize];
        }

        if let Some(node) = rover {
            if node.data.is_some() {
                node.data = Some(value);
                return true;
            }
            node.data = Some(value);
            node.use_count += 1;
            true
        } else {
            let mut node = Box::new(TrieNode::new());
            node.data = Some(value);
            node.use_count += 1;
            *rover = Some(node);
            true
        }
    }

    /// Removes a key from the trie.
    /// Returns `true` if the key was present and removed, `false` otherwise.
    pub fn remove(&mut self, key: &str) -> bool {
        let key_bytes = key.as_bytes();
        let mut path = VecDeque::new();
        let mut rover = &mut self.root_node;

        for &c in key_bytes {
            if let Some(node) = rover {
                path.push_back((c, node));
                rover = &mut node.next[c as usize];
            } else {
                return false;
            }
        }

        if let Some(node) = rover {
            if node.data.is_none() {
                return false;
            }
            node.data = None;
            self.cleanup_path(path);
            true
        } else {
            false
        }
    }

    /// Looks up a key in the trie and returns the associated value.
    /// Returns `None` if the key is not found.
    pub fn lookup(&self, key: &str) -> Option<&T> {
        let mut rover = &self.root_node;
        let key_bytes = key.as_bytes();

        for &c in key_bytes {
            if let Some(node) = rover {
                rover = &node.next[c as usize];
            } else {
                return None;
            }
        }

        rover.as_ref().and_then(|node| node.data.as_ref())
    }

    /// Returns the number of entries in the trie.
    pub fn num_entries(&self) -> usize {
        self.root_node.as_ref().map_or(0, |node| node.use_count)
    }

    /// Cleans up the path after a removal operation.
    fn cleanup_path(&mut self, mut path: VecDeque<(u8, &mut Box<TrieNode<T>>)>) {
        while let Some((c, parent)) = path.pop_back() {
            let node = &mut parent.next[c as usize];
            if let Some(node) = node {
                node.use_count -= 1;
                if node.use_count == 0 && node.data.is_none() {
                    *node = None;
                }
            }
        }
    }
}

#[cfg(test)]
mod tests {
    use super::*;

    #[test]
    fn test_insert_and_lookup() {
        let mut trie = Trie::new();
        assert!(trie.insert("key1", 10));
        assert!(trie.insert("key2", 20));
        assert_eq!(trie.lookup("key1"), Some(&10));
        assert_eq!(trie.lookup("key2"), Some(&20));
        assert_eq!(trie.lookup("key3"), None);
    }

    #[test]
    fn test_remove() {
        let mut trie = Trie::new();
        assert!(trie.insert("key1", 10));
        assert!(trie.insert("key2", 20));
        assert!(trie.remove("key1"));
        assert_eq!(trie.lookup("key1"), None);
        assert_eq!(trie.lookup("key2"), Some(&20));
    }

    #[test]
    fn test_num_entries() {
        let mut trie = Trie::new();
        assert_eq!(trie.num_entries(), 0);
        assert!(trie.insert("key1", 10));
        assert_eq!(trie.num_entries(), 1);
        assert!(trie.insert("key2", 20));
        assert_eq!(trie.num_entries(), 2);
        assert!(trie.remove("key1"));
        assert_eq!(trie.num_entries(), 1);
    }
}
